#!/usr/bin/python
'''
From an annotated edge file it pulls the correctly mapped edges.
'''
import sys
import os

# parameters.
node_file = sys.argv[1]
edge_file = sys.argv[2]
correct_file = sys.argv[3]
incorrect_file = sys.argv[4]

# Read in the node file.
print "Reading edge file."
fin = open(node_file, "rb")
lines = fin.readlines()
fin.close()

# Parse node file.
nodes = {}
print "Parsing node file."
for line in lines:
	# tokenize.
	tmp = line.strip().split("\t")
	
	# save id to chrom name.
	nodes[tmp[0]] = tmp

# Read in edge file.
print "Reading edge file."
fin = open(edge_file, "rb")
lines = fin.readlines()
fin.close()

# Parse edge file.
bad_pairs = []
good_pairs = []
print "Parsing edge file."
for line in lines:
	# tokenize.
	tmp = line.strip().split("\t")
	id1 = tmp[0]
	id2 = tmp[1]
	lstart1 = int(tmp[4])
	lstart2 = int(tmp[5])

	# sanity check.
	if id1 not in nodes:
		print "node error", id1
		sys.exit()	
	if id2 not in nodes:
		print "node error", id2
		sys.exit()

	# get global starts.
	gstart1 = int(nodes[id1][4]) + lstart1
	gstart2 = int(nodes[id2][4]) + lstart2
		
	# setup result to add.
	res = "%s\t%i\t%s\t%i\n" % (nodes[id1][1], lstart1, nodes[id2][1], lstart2)
		
	# Check for inter chromosome pairs.
	if nodes[id1][3] != nodes[id2][3]:
		bad_pairs.append(res)
		continue
		
	# Check for greater than 5 stddev.
	gstart1 = int(nodes[id1][4]) + lstart1
	gstart2 = int(nodes[id2][4]) + lstart2
	
	if abs(gstart1 - gstart2) > ((5 * 250) + 1395):
		bad_pairs.append(res)
		continue
	
	if abs(gstart1 - gstart2) < ((5 * 250) - 1395):
		bad_pairs.append(res)
		continue
		
	# Must be a good pair.
	good_pairs.append(res)
	
	
# Print out bad.
fout = open(incorrect_file, "wb")
for pair in bad_pairs:
	fout.write(pair)
fout.close()

# print out good.
fout = open(correct_file, "wb")
for pair in good_pairs:
	fout.write(pair)
fout.close()
